import cv2
import numpy as np
import os
import joblib
import gradio as gr
from gradio.components import Image  # 从这里导入 Image 组件

def load_best_model(model_path):
    return joblib.load(model_path)

def preprocess_image(image):
    img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
    img = cv2.resize(img, (32, 32))
    return img.reshape(1, -1)

def predict_image(input_image):
    model_path = os.path.join(os.getcwd(), "LGBMClassifier_best_model.pkl")
    best_classifier = load_best_model(model_path)
    preprocessed_img = preprocess_image(input_image)
    prediction = best_classifier.predict(preprocessed_img)
    return "dog" if prediction[0] == 1 else "cat"

# 创建 Gradio 界面
iface = gr.Interface(
    fn=predict_image,
    inputs=Image(),  # 使用从正确模块导入的 Image 组件
    outputs="text",
    title="猫狗分类器",
    description="上传一张图片，判断是猫还是狗。"
)

iface.launch()